import os
from traceback import print_tb
from jax import grad
import jax.numpy as jnp
from jax import jit
import time
import numpy as np
import numpy.random as npr
import jax
import jax.numpy as jnp
from jax import device_put
from jax import jit, grad, lax, random
from jax.example_libraries import optimizers
from jax.example_libraries import stax
from jax.example_libraries.stax import Dense, FanOut, Relu, Softplus, Sigmoid, FanInSum
from jax.nn import sigmoid
from functools import partial
from jax import vmap
from flax import linen as nn
from flax.training import train_state
from flax import struct
from jax import lax

from jax import tree_util
from jax.tree_util import tree_structure
from jax.tree_util import tree_flatten, tree_unflatten

import jax.experimental.sparse as sparse

import optax
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA

v1 = np.array([-i for i in range(3)])
v1 = v1.astype(float)
v2 = v1.copy()

# # 给 v2 的每一位施加一个随机扰动
# v2 += np.random.uniform(-20, 20, 100)

print(v1)
print(v2)

print("proj v1 on v2: ", np.dot(v1, v2) / np.linalg.norm(v1))

proj_err = np.abs(np.dot(v1, v2) / np.linalg.norm(v1) - np.linalg.norm(v1))
print("proj_err: ", proj_err)

dist_err = np.abs(np.linalg.norm(v1-v2) - np.linalg.norm(v1))
print("dist_err: ", dist_err)

# # 将 v1 和 v2 并列显示在两个子图中，绘制成bar chart
# fig, (ax1, ax2) = plt.subplots(1, 2)
# ax1.bar(np.arange(100), v1)
# ax2.bar(np.arange(100), v2)
# plt.show()

